package com.cn.ycy.filter;

import com.cn.ycy.filter.entity.CachedBodyOutputMessage;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.NettyWriteResponseFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractNameValueGatewayFilterFactory;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR;

/**
 * 修改响应
 */
@Component
public class SelfModifyResponseBodyGatewayFilterFactory extends AbstractNameValueGatewayFilterFactory {

  private static Logger logger = LoggerFactory.getLogger(SelfModifyResponseBodyGatewayFilterFactory.class);

  @Override
  public GatewayFilter apply(NameValueConfig config) {
    return new ModifyResponseGatewayFilter();
  }

  class ModifyResponseGatewayFilter implements GatewayFilter, Ordered {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
      return chain.filter(exchange.mutate().response(decorate(exchange)).build());
    }

    @SuppressWarnings("unchecked")
    private ServerHttpResponse decorate(ServerWebExchange exchange) {
      return new ServerHttpResponseDecorator(exchange.getResponse()) {

        @Override
        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {

//          Class inClass = config.getInClass();
//          Class outClass = config.getOutClass();

          String originalResponseContentType = exchange
                  .getAttribute(ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);
          HttpHeaders httpHeaders = new HttpHeaders();
          // explicitly add it in this way instead of
          // 'httpHeaders.setContentType(originalResponseContentType)'
          // this will prevent exception in case of using non-standard media
          // types like "Content-Type: image"
          httpHeaders.add(HttpHeaders.CONTENT_TYPE, originalResponseContentType);

          ClientResponse clientResponse = ClientResponse
                  .create(exchange.getResponse().getStatusCode())
                  .headers(headers -> headers.putAll(httpHeaders))
                  .body(Flux.from(body)).build();

          // TODO: flux or mono
          Mono modifiedBody = clientResponse.bodyToMono(String.class)
                  .flatMap(originalBody -> {
                    logger.info("response {}", originalBody);
                    //加密
                    return Mono.just("我是网关，我改写了响应");
                  });

          BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody,
                  String.class);
          CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(
                  exchange, exchange.getResponse().getHeaders());
          return bodyInserter.insert(outputMessage, new BodyInserterContext())
                  .then(Mono.defer(() -> {
                    Flux<DataBuffer> messageBody = outputMessage.getBody();
                    HttpHeaders headers = getDelegate().getHeaders();
                    if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING)) {
                      messageBody = messageBody.doOnNext(data -> headers
                              .setContentLength(data.readableByteCount()));
                    }
                    // TODO: fail if isStreamingMediaType?
                    return getDelegate().writeWith(messageBody);
                  }));
        }

        @Override
        public Mono<Void> writeAndFlushWith(
                Publisher<? extends Publisher<? extends DataBuffer>> body) {
          return writeWith(Flux.from(body).flatMapSequential(p -> p));
        }
      };
    }

    @Override
    public int getOrder() {
      return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 1;
    }

  }


}
